This is very cool @MasonProtter . Before reviewing, I just want to make sure that I've correctly understood what you've done in this PR. Is the below roughly correct?

"This PR introduces a generated function which returns the DerivedRule associated to a given signature. It adds backedges in the generated function to make sure that recompilation happens if a method is added in a later world age which would cause the rule to need to be re-derived. The advantage of this general approach is that the time it takes to get hold of a rule for a given signature is greatly reduced."

Yes @willtebbutt that's correct. However, there's also a catch with this approach, which gives it some downsides: it seems like an opaque closure created in a generated function cannot be inlined, which means that you hit allocations from actually calling this rule.

I have a branch where I did did this for the forward diff stuff, and the effect is easier to see there:

julia> function foo(f, x)
           rule = Mooncake.generated_build_frule(Tuple{typeof(f), Float64})
           rule(Dual(f, Mooncake.NoTangent()), Dual(x, 1.0))
       end;

julia> @btime foo(x -> sin(x + 1), 1.0)
  39.581 ns (3 allocations: 96 bytes)
Dual{Float64, Float64}(0.9092974268256817, -0.4161468365471424)

whereas if we compute the rule 'normally' the compiler is allowed to call it without overhead:

julia> function bar(rule, f, x)
           rule(Dual(f, Mooncake.NoTangent()), Dual(x, 1.0))
       end;

julia> let f = x -> sin(x + 1)
           rule = Mooncake.build_frule(Mooncake.get_interpreter(), Tuple{typeof(f), Float64})
           @btime bar($rule, $f, 1.0)
       end
  10.881 ns (0 allocations: 0 bytes)
Dual{Float64, Float64}(0.9092974268256817, -0.4161468365471424)

I'll investigate a bit at some point if this is possible to overcome.

Great, glad I've understood. Regarding the performance: it seems odd that this introduces allocations. Typically all of the operations associated to passing OpaqueClosures around and calling them are allocation-free, so I feel like it ought to be possible to get rid of these allocations.

Passing around opaque closures is cheap when the compiler is allowed to inline them, but if it's not, then they're expensive. Here's an example of that:

julia> function make_oc()
           Base.Experimental.@opaque (x::Int) -> sin(x - 1)
       end;

julia> @generated function make_oc_gen()
           Base.Experimental.@opaque (x::Int) -> sin(x - 1)
       end;

julia> @btime f(x) setup=begin
           x = 1
           f = make_oc()
       end
  3.566 ns (0 allocations: 0 bytes)
0.0

julia> @btime f(x) setup=begin
           x = 1
           f = make_oc_gen()
       end
  38.581 ns (2 allocations: 64 bytes)
0.0

You likely need to add __precompile__(false) statement somewhere to prohibit this package from being accessed via precompile. I’m often asked what things @generated is not allowed to do, and this is one of those things it is not allowed to do: using the precompile cache if any of its generated functions may return an OpaqueClosure method. Without that annotation, the package may cause incorrect or bad inference results which may poison any downstream consumer packages. This won’t generally be visible in tests, since tests don’t typically define downstream consumer packages, but it would be potentially a soundness and performance risk for downstream users if Mooncake fails to add that directive when making this addition of a generated function.

Thanks for letting us know! I would certainly not have thought about this.

So we'll likely need a different approach then rather than disabling precompilation (unless we stick the compile-time-generated version in a separate module that has precompilation disabled? But that's a pretty unsatisfying solution and probably quite brittle).

An alternative route we could take here would be to wait for v1.12, and use the mechanism introduced in JuliaLang/julia#56660 to try and do this. Of course, if JuliaLang/julia#56650 were to be resurrected that'd be much better for Mooncake. It seems like @Keno is still somewhat open to that design and having it complement 56660.

Maybe we can discuss with him sometime.

Closed in favour of #900